import pandas as pd
import logging
import numpy as np


def remove_which_cols_pearsonr(r, threshold=0.95, to_remove=2):
    '''
    当两两列的pearson r过大时，应该移除掉其中一列。
    这个函数返回应该移除哪些列。
    输入参数：
        r：格式：index是两层表示两列。有一列abs_pearson表示pearson r的绝对值。
        to_remove：是移除第一层的列还是第二层的列。
        threshold：>=这个threshold的才会考虑移除之一。
    输出参数：
        一个list，表示需要移除的columns。
    '''
    r = r[r.abs_pearson >= threshold]
    cols_list1 = r.index.get_level_values(0).tolist()
    cols_list2 = r.index.get_level_values(1).tolist()
    all_remove = set()

    for i in range(len(cols_list1)):
        if (cols_list1[i] not in all_remove) and(cols_list2[i] not in all_remove):  # 必须要去除其中之一
            if to_remove == 1:
                all_remove.add(cols_list1[i])
            else:
                all_remove.add(cols_list2[i])
    all_remove = list(all_remove)
    return all_remove


def print_full(x):
    pd.set_option('display.max_rows', len(x))
    print(x)
    pd.reset_option('display.max_rows')


def drop_if_possible(df, to_drop, inplace=True):
    assert (df.columns.duplicated() == False).all()
    real_to_drop = []
    for col in to_drop:
        if col in df.columns:
            real_to_drop.append(col)
    logging.warning('df previous cols: %d origin to_drop cols: %d will drop cols: %d now df cols: %d' % (
        len(df.columns),
        len(to_drop),
        len(real_to_drop),
        len(df.columns) - len(real_to_drop)
    )
    )
    if inplace:
        df.drop(real_to_drop, axis=1, inplace=True)
    else:
        return df.drop(real_to_drop, axis=1)


def iter_n_rows(X, n):
    '''
    生成器，X是ndarray，每次返回n行直到结束
    '''
    len_row = X.shape[0]
    start = 0
    while start < len_row:
        if start + n >= len_row:
            yield X[start:]
            break
        else:
            yield X[start:start + n]
            start += n


def iter_n_rows_list(X_list, n):
    '''
    生成器，X_list，y, 每次返回n行并拼接
    '''
    len_row = X_list[0].shape[0]
    start = 0
    while start < len_row:
        if start + n >= len_row:
            yield np.concatenate([X[start:] for X in X_list], axis=1)
            break
        else:
            yield np.concatenate([X[start:start + n] for X in X_list], axis=1)
            start += n


def get_cols_index(df, col_name_list):
    index = []
    assert (df.columns.duplicated() == False).all()
    for col in col_name_list:
        index.append((df.columns == col).nonzero()[0][0])
    return index
